#!/usr/bin/env python3
"""
数据库迁移脚本：为ComplementaryFoodRecord表添加acceptance_level字段
"""

from flask import Flask
from db import db
from config import Config
import pymysql

app = Flask(__name__)
app.config.from_object(Config)
db.init_app(app)

def migrate_database():
    """执行数据库迁移"""
    with app.app_context():
        try:
            # 连接到MySQL数据库
            connection = pymysql.connect(
                host='localhost',
                user='root',
                password='qwer1234',
                database='mamaassist',
                charset='utf8mb4'
            )
            
            with connection.cursor() as cursor:
                # 检查字段是否已存在
                cursor.execute("""
                    SELECT COUNT(*) 
                    FROM INFORMATION_SCHEMA.COLUMNS 
                    WHERE TABLE_SCHEMA = 'mamaassist' 
                    AND TABLE_NAME = 'complementary_food_record' 
                    AND COLUMN_NAME = 'acceptance_level'
                """)
                
                field_exists = cursor.fetchone()[0] > 0
                
                if not field_exists:
                    print("添加 acceptance_level 字段到 complementary_food_record 表...")
                    cursor.execute("""
                        ALTER TABLE complementary_food_record 
                        ADD COLUMN acceptance_level VARCHAR(20) DEFAULT ''
                    """)
                    print("字段添加成功！")
                else:
                    print("acceptance_level 字段已存在，跳过添加。")
                
                # 检查并移除 complementary_food_detail 表中的 acceptance_level 字段
                cursor.execute("""
                    SELECT COUNT(*) 
                    FROM INFORMATION_SCHEMA.COLUMNS 
                    WHERE TABLE_SCHEMA = 'mamaassist' 
                    AND TABLE_NAME = 'complementary_food_detail' 
                    AND COLUMN_NAME = 'acceptance_level'
                """)
                
                detail_field_exists = cursor.fetchone()[0] > 0
                
                if detail_field_exists:
                    print("从 complementary_food_detail 表移除 acceptance_level 字段...")
                    cursor.execute("""
                        ALTER TABLE complementary_food_detail 
                        DROP COLUMN acceptance_level
                    """)
                    print("字段移除成功！")
                else:
                    print("complementary_food_detail 表中没有 acceptance_level 字段，跳过移除。")
            
            connection.commit()
            print("数据库迁移完成！")
            
        except Exception as e:
            print(f"迁移失败: {e}")
            connection.rollback()
        finally:
            connection.close()

if __name__ == '__main__':
    migrate_database()